package datastruct;

import org.junit.Test;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

public class SubArraySum560 {

    @Test
    public void test() {
        subarraySum(new int[]{1, 1, 1}, 1);
        subarraySum(new int[]{1, 1, 1}, 2);
        subarraySum(new int[]{1, 2, 3}, 2);
        subarraySum(new int[]{1, 2, 2}, 4);

        subarraySum2(new int[]{1, 1, 1}, 1);
        subarraySum2(new int[]{1, 1, 1}, 2);
        subarraySum2(new int[]{1, 2, 3}, 2);
        subarraySum2(new int[]{1, 2, 2}, 4);
    }

    public int subarraySum(int[] nums, int k) {
        // prefixSum[i+1] -> nums[i]
        int[] prefixSum = new int[nums.length+1];
        prefixSum[0] = 0;
        for (int i = 0; i < nums.length; i++) {
            prefixSum[i+1] = prefixSum[i] + nums[i];
        }

        int count = 0;
        for (int i = 0; i < nums.length; i++) {
            for (int j = i+1; j <= nums.length; j++) {
                if (prefixSum[j] - prefixSum[i] == k) {
                    count++;
                }
            }
        }
        System.out.format("nums: %s, k: %d, count: %d\n", Arrays.toString(nums), k, count);
        return count;
    }

    public int subarraySum2(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>(nums.length * 3 / 2);
        int prefixSum = 0;
        int count = 0;
        
        map.put(0, 1);
        for (int i = 0; i < nums.length; i++) {
            prefixSum += nums[i];
            count += map.getOrDefault(prefixSum-k, 0);
            map.put(prefixSum, map.getOrDefault(prefixSum, 0)+1);
        }
        System.out.format("nums: %s, k: %d, count: %d\n", Arrays.toString(nums), k, count);
        return count;
    }
}
